import os

import grpc
import service.pb.model_server_pb2 as pb
from service.pb.model_server_pb2_grpc import ModelServerStub
from service.tools import save_chunks_to_file

class ModelServerClient:
    def __init__(self, ip, port=50051):
        self.channel = grpc.insecure_channel(f'{ip}:{port}')
        self.stub = ModelServerStub(self.channel)

    def dataset_download(self, task, mode, dataset_name, output_file=None):
        request = pb.DatasetDownloadRequest(task=task, mode=mode, dataset_name=dataset_name)
        chunks = self.stub.dataset_download(request)
        if output_file is None:
            output_file = f'{task}_{mode}_{dataset_name}.json'
        save_chunks_to_file(chunks, output_file)
        assert os.path.exists(output_file)

    def close(self):
        self.channel.close()
